import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 去警告
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # 使用CPU

import gzip
import pickle
import tensorflow as tf
from tensorflow import keras


# # 从网络中下载数据集代码
# from pathlib import Path
# import requests
# DATA_PATH = Path("data")
# PATH = DATA_PATH / "mnist"
# PATH.mkdir(parents=True, exist_ok=True)
# URL = "http://..."
# FILENAME = "mnist.pkl.gz"
# if not (PATH / FILENAME).exists():
#     content = requests.get(URL + FILENAME).content
#     (PATH / FILENAME).open("wb").write(content)

# # 加载数据集
# from tensorflow import keras
# fashion_mnist = keras.datasets.fashion_mnist
# (train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()



if __name__ == "__main__":
    # 1. 读取数据
    file_dir = '../resources/p02_deep_learning_tensorFlow/mnist/mnist.pkl.gz'
    with gzip.open(file_dir, "rb") as f:
        ((x_train, y_train), (x_valid, y_valid), _) = pickle.load(f, encoding="latin-1")

    # plt.imshow(x_train[0].reshape((28, 28)), cmap="gray")  # 展示数据（28*28*1）的灰度图，需要预处理成 28*28 个特征

    # 2.数据处理
    train = tf.data.Dataset.from_tensor_slices((x_train, y_train))
    train = train.batch(32) # 每32个张量数据构成一个张量数据
    train = train.repeat()

    valid = tf.data.Dataset.from_tensor_slices((x_valid, y_valid))
    valid = valid.batch(32) # 每32个张量数据构成一个张量数据
    valid = valid.repeat()

    # 3. 特征工程
    # 估计器：基于 Keras 构建网络模型
    # 前向传播：按顺序构造网络模型
    model = keras.Sequential([
        # keras.layers.Flatten(input_shape=(28, 28)),  #拉长数据,如果进行了 data 则不需要
        keras.layers.Dense(32, activation='relu'),  # activation：激活函数的选择，一般是relu
        keras.layers.Dense(32, activation='relu'),
        keras.layers.Dense(10, activation='softmax')
    ])

    # 反向传播：无连续目标值（监督学习） - 分类问题 - 优化器 - 指定优化器、损失函数
    model.compile(optimizer=tf.keras.optimizers.Adam(0.001),
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(),   # 注意：不同任务的损失函数 loss 不同
                  metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])  # metrics：展示准确率
    # model.fit(x_train, y_train, epochs=5, batch_size=64, validation_data=(x_valid, y_valid))  # x , 迭代5次，一次迭代64样本（越大越好）, y
    model.fit(train, epochs=5, steps_per_epoch=100, validation_data=valid, validation_steps=100)  # x , 迭代5便，便迭代100次，一次迭代64样本（越大越好）, y ，

    # 4.评估模型
    # 损失值。ACC
    test_loss, test_acc = model.evaluate(x_valid, y_valid, verbose=2)
    # 预测值
    y_pred = model.predict(x_valid)


    ## 5.保存模型
    ## 5.1 直接保存成 1 个文件
    ## 保存模型
    # model.save("../resources/p02_deep_learning_tensorFlow/mnist/model/mnist_model.h5")
    ## 加载模型
    # model = keras.models.load_model("../resources/p02_deep_learning_tensorFlow/mnist/model/mnist_model.h5")

    # # 5.2 保存成 json
    # # 保存模型
    # config = model.to_json()
    # with open('config.json', 'w') as json:
    #     json.write(config)
    # # 加载模型
    # json_config = 'config.json'
    # model = keras.models.model_from_json(json_config)
    # weights = model.get_weights()